package com.zxt.spring_ai_demo.controller;

import com.zxt.spring_ai_demo.tool.DateTimeTools;
import com.zxt.spring_ai_demo.tool.WeatherService;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.content.Media;
import org.springframework.ai.image.ImagePrompt;
import org.springframework.ai.image.ImageResponse;
import org.springframework.ai.model.tool.ToolCallingChatOptions;
import org.springframework.ai.openai.*;
import org.springframework.ai.openai.api.OpenAiAudioApi;
import org.springframework.ai.openai.audio.speech.SpeechPrompt;
import org.springframework.ai.openai.audio.speech.SpeechResponse;
import org.springframework.ai.support.ToolCallbacks;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.tool.function.FunctionToolCallback;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.io.ClassPathResource;
import org.springframework.util.MimeTypeUtils;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.io.FileOutputStream;
import java.io.IOException;

/**
 * @Author Mr.zeng
 * @Date 2025/8/13 20:48
 */
@RestController
@RequestMapping("/openAI")
public class OpenAIController {
    // 聊天模型
    @Autowired
    private OpenAiChatModel openAiChatModel;
    // 生成图片的模型
    @Autowired
    private OpenAiImageModel openaiImageModel;
    // 生成音频模型
    @Autowired
    private OpenAiAudioSpeechModel openAiAudioSpeechModel;

    @GetMapping("/chat")
    public String chat(String userInput) {
        String result = openAiChatModel.call(userInput);
        return result;
    }

    @GetMapping(value = "/streamChat", produces = "text/plain;charset=UTF-8")
    public Flux<String> streamChat(String userInput) {
        Flux<String> result = openAiChatModel.stream(userInput);
        return result;
    }

    // 根据用户输入的描述生成图片
    @GetMapping("/image")
    public String image(String userInput) {
        ImageResponse response = openaiImageModel.call(
                new ImagePrompt(userInput,
                        OpenAiImageOptions.builder()
                                .quality("hd")  // 质量，高清
                                .N(1)   // The number of images to generate. Must be between 1 and 10. For dall-e-3, only n=1 is supported.
                                .height(1024)
                                .width(1024).build())

        );
        String imageUrl = response.getResult().getOutput().getUrl();
        String result = "<img src=" + imageUrl + "/>";
        return result;
    }

    // 文字转音频
    @GetMapping("/audio")
    public String audio(String userInput) throws IOException {
        OpenAiAudioSpeechOptions speechOptions = OpenAiAudioSpeechOptions.builder()
                .model("tts-1") // 使用哪个模型
                .voice(OpenAiAudioApi.SpeechRequest.Voice.NOVA)    // 设置音色
                .responseFormat(OpenAiAudioApi.SpeechRequest.AudioResponseFormat.MP3)   // 音频格式
                .speed(1.0f)    // 音频速度
                .build();

        SpeechPrompt speechPrompt = new SpeechPrompt(userInput, speechOptions);
        SpeechResponse response = openAiAudioSpeechModel.call(speechPrompt);
        byte[] bytes = response.getResult().getOutput();
        // 将音频信息字节数组写到本地文件中
        FileOutputStream fileOutputStream = new FileOutputStream("user.mp3");
        fileOutputStream.write(bytes);
        fileOutputStream.close();
        return "音频生成成功";
    }

    // 多模态输入
    @GetMapping("/multiModel")
    public String multiModel(String userInput) throws IOException {
        // 准备图片
        ClassPathResource resource = new ClassPathResource("/static/zxt.jpeg");
        // 图片封装为Media对象
        Media media = new Media(MimeTypeUtils.IMAGE_JPEG, resource);
        // 封装成message对象，必须使用user角色
        UserMessage userMessage = UserMessage.builder()
                .text(userInput)
                .media(media)
                .build();

        ChatResponse response = this.openAiChatModel.call(new Prompt(userMessage));
        return response.getResult().getOutput().getText();
    }

    // 函数调用，tool方式
    @GetMapping("/tool")
    public String tool(String userInput) {
        ToolCallback[] dateTimeTools = ToolCallbacks.from(new DateTimeTools());
        // 创建聊天配置对象
        ChatOptions chatOptions = ToolCallingChatOptions.builder()
                .toolCallbacks(dateTimeTools)
                .build();
        Prompt prompt = new Prompt(userInput, chatOptions);
        ChatResponse chatResponse = openAiChatModel.call(prompt);
        String result = chatResponse.getResult().getOutput().getText();
        return result;
    }

    // 函数调用，Fuction方式
    @GetMapping("/function")
    public String function(String userInput) {
        // 使用自定义函数类创建一个toolCallback函数回调
        ToolCallback toolCallback = FunctionToolCallback
                .builder("currentWeather", new WeatherService())
                .description("查询天气")
                .inputType(WeatherService.WeatherRequest.class)
                .build();
        ChatOptions chatOptions  = ToolCallingChatOptions.builder().toolCallbacks(toolCallback).build();
        Prompt prompt = new Prompt(userInput, chatOptions);
        ChatResponse chatResponse = openAiChatModel.call(prompt);
        String result = chatResponse.getResult().getOutput().getText();
        return result;
    }
}
